package com.iakuil.app.common.tp;

import cn.hutool.core.thread.NamedThreadFactory;
import cn.hutool.core.thread.RejectPolicy;
import com.iakuil.app.common.core.constant.SysConstant;
import com.iakuil.app.common.core.util.ObjectUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.support.BeanDefinitionBuilder;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.boot.context.properties.bind.Bindable;
import org.springframework.boot.context.properties.bind.Binder;
import org.springframework.context.EnvironmentAware;
import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
import org.springframework.core.ResolvableType;
import org.springframework.core.env.Environment;
import org.springframework.core.type.AnnotationMetadata;

import java.util.Arrays;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * 自定义线程池注册
 *
 * <p>配置示例：
 * <pre>{@code
 * app:
 *   tp:
 *     enabled: true
 *     pools:
 *       -  name: test
 *          isDaemon: false
 *          corePoolSize: 2
 *          maxPoolSize: 4
 *          keepAliveSeconds: 30
 *          queueCapacity: 10
 *          awaitTerminationMillis: 60
 * }</pre>
 *
 * @author Kai
 */
@Slf4j
public class CustomThreadPoolRegistrar implements ImportBeanDefinitionRegistrar, EnvironmentAware {
    private Environment environment;

    @Override
    public void registerBeanDefinitions(AnnotationMetadata importingClassMetadata, BeanDefinitionRegistry registry) {
        // 绑定资源
        ThreadPoolProperties tpProperties = new ThreadPoolProperties();
        Binder binder = Binder.get(environment);
        ResolvableType resolvableType = ResolvableType.forClass(ThreadPoolProperties.class);
        Bindable<Object> bindable = Bindable.of(resolvableType).withExistingValue(tpProperties);
        binder.bind(SysConstant.APP_SYMBOL + ".tp", bindable);

        if (!tpProperties.getEnabled()) {
            return;
        }

        ThreadPoolProperties.Prop[] executors = tpProperties.getPools();
        if (ObjectUtils.isEmptyOrNull(executors)) {
            log.warn("未检测到自定义线程池配置");
            return;
        }

        // 依次注册
        Arrays.stream(executors).forEach((executorProp) -> {
            BeanDefinitionBuilder builder = BeanDefinitionBuilder.genericBeanDefinition(ThreadPoolExecutor.class);
            for (Object arg : assembleArgs(executorProp)) {
                builder.addConstructorArgValue(arg);
            }
            registry.registerBeanDefinition(executorProp.getName(), builder.getBeanDefinition());
        });
    }

    @Override
    public void setEnvironment(Environment environment) {
        this.environment = environment;
    }

    private Object[] assembleArgs(ThreadPoolProperties.Prop prop) {
        String poolName = prop.getName();
        return new Object[]{
            prop.getCorePoolSize(),
            prop.getMaxPoolSize(),
            prop.getKeepAliveSeconds(),
            TimeUnit.SECONDS,
            new ArrayBlockingQueue<>(prop.getQueueCapacity()),
            new NamedThreadFactory(poolName + "-thread-", prop.isDaemon()),
            // 全部采用ABORT拒绝策略
            RejectPolicy.ABORT.getValue()
        };
    }
}
